# 导入必要的库
import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import PolynomialFeatures  # 多项式特征生成
from sklearn.linear_model import LinearRegression  # 线性回归模型
from sklearn.metrics import mean_squared_error  # 均方误差评估

# 初始化设置
plt.rcParams.update({
    'font.size': 14,  # 全局字体大小
    'axes.titlesize': 16,  # 标题字体大小
    'axes.labelsize': 12,  # 坐标轴标签字体大小
    'xtick.labelsize': 12,  # x轴刻度字体
    'ytick.labelsize': 12,  # y轴刻度字体
    'legend.fontsize': 12  # 图例字体
})
plt.rcParams['font.sans-serif'] = ['KaiTi']  # 中文字体
plt.rcParams['mathtext.fontset'] = 'stix'  # 数学字体
plt.rcParams['axes.unicode_minus'] = False

# 设置随机种子保证结果可复现
np.random.seed(42)

# 生成模拟数据
m = 200  # 样本数量
X = 6 * np.random.rand(m, 1) - 3  # 生成-3到3之间的随机特征
y = 2 + X + 0.5 * X ** 2 + np.random.randn(m, 1)  # 规范顺序：常数项 + 一次项 + 二次项 + 噪声
# plt.plot(X, y, 'k--', linewidth=2, label='真实模型')

# 绘制原始数据散点图
plt.figure(figsize=(10, 6))
plt.plot(X, y, 'b.', label='原始数据')

# 划分训练集和测试集 (80%训练，20%测试)
X_train = X[:150]
y_train = y[:150]
X_test = X[150:]
y_test = y[150:]



# 定义不同多项式阶数的样式和标签
d = {
    1: ('g-', '线性回归'),
    2: ('r+', '二次多项式'),
    10: ('y*', '十次多项式')
}

print("=" * 50)
print("多项式回归模型比较".center(40))
print("=" * 50)

for degree in d:
    # 1. 生成多项式特征
    poly_features = PolynomialFeatures(degree=degree, include_bias=True)
    X_poly_train = poly_features.fit_transform(X_train)  # 训练集拟合并转换
    X_poly_test = poly_features.transform(X_test)  # 测试集仅转换

    # 2. 训练线性回归模型
    lin_reg = LinearRegression(fit_intercept=True)
    lin_reg.fit(X_poly_train, y_train)

    # 3. 预测结果
    y_train_predict = lin_reg.predict(X_poly_train)
    y_test_predict = lin_reg.predict(X_poly_test)

    # 4. 绘制拟合曲线
    plt.plot(X_poly_train[:, 1], y_train_predict, d[degree][0], label=d[degree][1])


    # 5. 计算并打印评估指标
    train_mse = mean_squared_error(y_train, y_train_predict)
    test_mse = mean_squared_error(y_test, y_test_predict)

    print(f"多项式阶数: {degree}")

    print(f"训练集维度: {X_poly_train.shape} (原始特征维度: {X_train.shape})")
    print(f"测试集维度: {X_poly_test.shape} (原始特征维度: {X_test.shape})")
    print(f"训练集原始样本 X: {X_train[0].round(4)}")
    print(f"训练集升维样本 X_poly: {X_poly_train[0].round(4)}")
    print(f"训练集真实值 y: {y_train[0].round(4)}")
    print(f"训练集预测值 y: {y_train_predict[0].round(4)}")


    print(f"训练集MSE: {train_mse:.4f}")
    print(f"测试集MSE: {test_mse:.4f}")
    print(f"模型截距: {lin_reg.intercept_.round(4)}")
    print(f"模型系数: {np.round(lin_reg.coef_, 4)}")

    print("=" * 50)

# 绘制图例和网格线
plt.xlabel("特征值 X")
plt.ylabel("目标值 y")
plt.title("不同阶数多项式回归比较")
plt.legend()
plt.grid(True)  # 添加网格线
plt.show()
